#coding=utf-8
import tensorflow as tf
import numpy as np
import vector

FLAGS = tf.app.flags.FLAGS
tf.app.flags.DEFINE_string("train_dir", "../train_dir/*", "train dir")

tf.app.flags.DEFINE_string("head_user_list", "../vocab.txt", "head user list")
tf.app.flags.DEFINE_string("user_list", "./user_list.txt", "user list")

tf.app.flags.DEFINE_string("head_model_dir", "../model_ckpt", "head user model dir")
tf.app.flags.DEFINE_string("model_ckpt", "./model_ckpt/user-model.ckpt", "all user model ckpt")
tf.app.flags.DEFINE_integer("embedding_size", 128, "embedding size")


def load_vocab(vocab_file):
	words = []
	with open(vocab_file, 'r') as vocab_stream:
		for line in vocab_stream:
			word = line.replace('\n', '')
			words.append(word)
	return words

def main(argv):
	vec = vector.vector(embedding_size=FLAGS.embedding_size,
						head_user_list=load_vocab(FLAGS.head_user_list),
						user_list=load_vocab(FLAGS.user_list),
						head_model_dir=FLAGS.head_model_dir,
						model_ckpt=FLAGS.model_ckpt,
						train_dir=FLAGS.train_dir)
	vec.start()	
	vec.save_model()
	vec.print_model()

if __name__ == '__main__':
	tf.logging.set_verbosity(tf.logging.INFO)
	tf.app.run(main)
